import os
import json
import requests
import queue
import threading
from concurrent.futures import ThreadPoolExecutor
from argparse import ArgumentParser


class BulkProductCreator:
    def __init__(self, base_url, image_dir):
        self.base_url = base_url
        self.image_dir = image_dir
        self.success_count = 0
        self.count_lock = threading.Lock()  # 定义计数锁

    def login_user(self, username, password):
        """用户登录获取token"""
        url = f"{self.base_url}/api/users/login/"
        try:
            response = requests.post(url, json={"username": username, "password": password}, timeout=10)
            if response.status_code == 200:
                return response.json().get('access')
            print(f"[ERROR] 登录失败 {username}: {response.text}")
            return None
        except Exception as e:
            print(f"[ERROR] 登录异常 {username}: {str(e)}")
            return None

    def _create_request(self, user_token, product_data, image_paths):
        """构建并发送创建请求"""
        data = {
            'name': product_data['name'],
            'price': str(product_data['price']),
            'category': str(product_data['category']),
            'description': product_data['description'],
            'delivery_method': str(product_data['delivery_method']),
            'delivery_location': product_data['delivery_location'],
            'action': 'publish'
        }

        with requests.Session() as session:
            files = []
            try:
                for path in image_paths:
                    try:
                        f = open(path, 'rb')
                        files.append(('images', (os.path.basename(path), f, 'image/jpeg')))
                    except Exception as e:
                        print(f"[ERROR] 打开图片文件失败 {path}: {str(e)}")
                response = session.post(
                    f"{self.base_url}/api/products/seller/",
                    data=data,
                    files=files,
                    headers={
                        "Authorization": f"Bearer {user_token}",
                        "Accept": "application/json"
                    },
                    timeout=20
                )
            finally:
                for _, (_, f, _) in files:
                    f.close()

            return response

    def create_product(self, user_token, product_data, image_paths, max_retries=3):
        """带重试机制的商品创建（添加计数）"""
        for attempt in range(max_retries):
            try:
                response = self._create_request(user_token, product_data, image_paths)
                if response.status_code == 201:
                    with self.count_lock:  # 线程安全计数
                        self.success_count += 1
                    return True, ""
                return False, f"状态码 {response.status_code}: {response.text}"
            except Exception as e:
                if attempt == max_retries - 1:
                    return False, f"最终失败: {str(e)}"
                print(f"重试中 ({attempt + 1}/{max_retries}): {str(e)}")

    def worker(self, user_info, task_queue):
        """改进的工作线程，添加资源清理"""
        try:
            token = self.login_user(user_info['username'], user_info['password'])
            if not token:
                print(f"[DEBUG] 用户 {user_info['username']} 登录失败")
                return

            print(f"[INFO] 用户 {user_info['username']} 就绪")
            while True:
                try:
                    product_data, image_paths = task_queue.get_nowait()
                    print(f"[DEBUG] 用户 {user_info['username']} 正在处理: {product_data['name']}")
                    success, msg = self.create_product(token, product_data, image_paths)

                    if success:
                        print(f"[SUCCESS] {user_info['username']} 创建 {product_data['name']}")
                    else:
                        print(f"[ERROR] {user_info['username']} 失败: {msg}")

                    task_queue.task_done()
                except queue.Empty:
                    print(f"[DEBUG] 用户 {user_info['username']} 任务队列已空")
                    break
        finally:
            # 确保线程结束时释放资源
            print(f"[INFO] 用户 {user_info['username']} 工作完成")


def load_json_file(filepath):
    with open(filepath, 'r', encoding='utf-8') as f:
        return json.load(f)


def scan_images(image_dir):
    image_files = []
    if not os.path.exists(image_dir):
        print(f"[ERROR] 图片目录不存在: {image_dir}")
        return []
    for root, _, files in os.walk(image_dir):
        for file in files:
            if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                image_files.append(os.path.join(root, file))
    return [image_files[i:i + 2] for i in range(0, len(image_files), 2)]


def main():
    parser = ArgumentParser(description='高效商品发布工具')
    parser.add_argument('--userfile', required=True, help='用户凭证JSON文件')
    parser.add_argument('--productfile', required=True, help='商品数据JSON文件')
    parser.add_argument('--imagedir', required=True, help='商品图片目录')
    parser.add_argument('--baseurl', default='http://localhost:8000', help='API基础URL')
    parser.add_argument('--threads', type=int, default=5, help='并发工作线程数')

    args = parser.parse_args()
    users = load_json_file(args.userfile)
    # 改进任务分配：确保每个用户获得均衡的任务量
    task_queue = queue.Queue()
    products = load_json_file(args.productfile)
    images = scan_images(args.imagedir)
    print(f"[DEBUG] 加载商品数量: {len(products)}")
    print(f"[DEBUG] 加载图片组数: {len(images)}")
    # 按用户分组分配任务
    user_count = min(len(users), args.threads)
    tasks_per_user = len(products) // user_count
    print(f"[DEBUG] 用户数量: {user_count}")
    print(f"[DEBUG] 每个用户分配任务数: {tasks_per_user}")
    for i, user in enumerate(users[:user_count]):
        start_idx = i * tasks_per_user
        end_idx = (i + 1) * tasks_per_user if i != user_count - 1 else len(products)
        print(f"[DEBUG] 用户 {user['username']} 分配商品: {start_idx}-{end_idx}")
        for product, img_group in zip(products[start_idx:end_idx], images[start_idx:end_idx]):
            task_queue.put((product, img_group))
    print(f"[DEBUG] 任务队列大小: {task_queue.qsize()}")
    # 初始化创建器
    creator = BulkProductCreator(args.baseurl, args.imagedir)

    # 使用线程池执行
    with ThreadPoolExecutor(max_workers=args.threads) as executor:
        # 每个用户分配一个线程
        futures = [executor.submit(creator.worker, user, task_queue)
                   for user in users[:args.threads]]

        # 等待所有任务完成
        task_queue.join()

        # 清理线程
        for future in futures:
            future.cancel()

    print(f"[STATS] 总成功创建商品数: {creator.success_count}")  # 新增统计输出


if __name__ == '__main__':
    main()

    """
    python products/tests.py \
    --userfile ./test/user.json \
    --productfile ./test/products.json \
    --imagedir "C:\\Users\\ASUS\\Pictures\\testImages" \
    --threads 10
    """